import json
import google.protobuf.json_format as json_format
import coin.exchange.base.kr_rest.private_client_base as privcb
import coin.exchange.huobi_futures_swap.kr_rest.native_private_client as nprivc
from coin.exchange.huobi.kr_rest.currency import HuobiCurrency
from coin.exchange.huobi_futures.kr_rest.futures_product import HuobiFuturesProduct
from coin.exchange.huobi_futures_swap.kr_rest.product_info import (huobi_futures_swap_product_info,
                                                                   huobi_futures_linear_swap_product_info)
from coin2.exchange.order.system_pb2 import OrderSystemConfig
from coin2.exchange.symbology_pb2 import ProductConfig
from cc.appcoin2.support.order.py_parser import PyOrderParser

from coin.proto.coin_query_pb2 import (AccountBalance,
                                       CurrencyBalance,
                                       AccountPosition,
                                       ProductPosition)


class HuobiSwapFuturesPrivateParser(object):
  @staticmethod
  def parse_usdt_balance(update_msg, update_msg_cross):
    currency = HuobiCurrency.FromStrNativeCurrency('USDT')
    currency_balance = CurrencyBalance(currency=currency.currency,
                                       currency_native=currency.native_currency,
                                       total=0,
                                       hold=0,
                                       available=0)
    for balance_info in update_msg.each_balance:
      currency_balance.total += balance_info.total
      currency_balance.hold += balance_info.hold
      currency_balance.available += balance_info.available

    assert update_msg_cross['status'] == 'ok', update_msg_cross
    assert len(update_msg_cross['data']) == 1
    assert update_msg_cross['data'][0]['margin_account'] == 'USDT'
    assert update_msg_cross['data'][0]['margin_asset'] == 'USDT'
    balance_info = update_msg_cross['data'][0]
    currency_balance.total += float(balance_info['margin_static'])
    currency_balance.hold += float(balance_info['margin_position'])
    currency_balance.available += float(balance_info['margin_static']) - float(balance_info['margin_position'])

    return AccountBalance(
        exchange='Huobi',
        each_balance=[currency_balance],
    )

  @staticmethod
  def parse_balance(update_msg, update_msg_usdt, update_msg_cross_usdt):
    assert update_msg['status'] == 'ok'
    currency_balance_list = []
    update_msg_usdt = HuobiSwapFuturesPrivateParser.parse_usdt_balance(
        update_msg_usdt, update_msg_cross_usdt)
    if update_msg_usdt is not None:
      assert len(update_msg_usdt.each_balance) == 1
      usdt_currency_balance = update_msg_usdt.each_balance[0]
    else:
      usdt_currency_balance = None
    for balance_info in update_msg['data']:
      native_currency = balance_info['symbol']
      currency = HuobiCurrency.FromStrNativeCurrency(native_currency)
      if currency is None:
        continue

      unrealized_pnl = float(balance_info['profit_unreal'])
      total = float(balance_info['margin_balance'])
      hold = float(balance_info['margin_frozen'])
      available = float(balance_info['margin_available'])
      currency_balance = CurrencyBalance(currency=currency.currency,
                                         currency_native=currency.native_currency,
                                         total=total,
                                         hold=hold,
                                         available=available)
      if currency.currency == 'USDT' and usdt_currency_balance is not None:
        currency_balance.total += usdt_currency_balance.total
        currency_balance.hold += usdt_currency_balance.hold
        currency_balance.available += usdt_currency_balance.available
        usdt_currency_balance = None
      currency_balance_list.append(currency_balance)
    if usdt_currency_balance is not None:
      currency_balance_list.append(usdt_currency_balance)

    return AccountBalance(
        exchange='Huobi',
        each_balance=currency_balance_list,
    )

  @staticmethod
  def parse_account_position(update_msg):
    assert update_msg['status'] == 'ok', update_msg
    product_positions = {}
    for pos_info in update_msg['data']:
      native_symbol = pos_info['contract_code']
      product = HuobiFuturesProduct.FromStrNativeProduct(native_symbol)
      position = product_positions.get(native_symbol)
      if position is None:
        position = ProductPosition(symbol=product.symbol)
        product_positions[native_symbol] = position
      if pos_info['direction'] == 'buy':
        position.long_position = float(pos_info['volume'])
        position.available_long_position = float(pos_info['available'])
      elif pos_info['direction'] == 'sell':
        position.short_position = -float(pos_info['volume'])
        position.available_short_position = -float(pos_info['available'])
      else:
        raise ValueError('Invalid direction!')

    for position in product_positions.values():
      position.net_position = position.long_position + position.short_position

    return AccountPosition(exchange='Huobi', each_position=product_positions.values())


class HuobiLinearSwapFuturesPrivateClient(privcb.PrivateClientBase):
  ProductType = HuobiFuturesProduct

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.HuobiSwapNativePrivateClient(key_file=key_file)
    linear_swap_pi_bundle = huobi_futures_linear_swap_product_info._product_infos
    assert linear_swap_pi_bundle.HasField('mea')
    self._linear_parser = PyOrderParser.from_mea_str(linear_swap_pi_bundle.mea)
    linear_swap_oss_config = OrderSystemConfig(
      mea=linear_swap_pi_bundle.mea,
      products=ProductConfig(norms=[pi.symbol for pi in linear_swap_pi_bundle.product_infos]))
    self._linear_parser.add_product(linear_swap_oss_config.SerializeToString())

  def query_account_balance_impl(self):
    update_usdt = self.nprivc.query_usdt_account_info()
    update_cross_usdt = self.nprivc.query_cross_usdt_account_info()

    self._linear_parser.set_user_cross_mode(False)
    msg_usdt = self._linear_parser.parse_account_balance(json.dumps(update_usdt.msg))
    self._linear_parser.set_user_cross_mode(True)
    msg_cross_usdt = self._linear_parser.parse_account_balance(json.dumps(update_cross_usdt.msg))

    balance_list = []
    for msg in (msg_usdt, msg_cross_usdt):
      account_balance = AccountBalance()
      json_format.Parse(msg, account_balance)
      assert account_balance.HasField('market_type')
      assert account_balance.HasField('exchange')
      balance_list.append(account_balance)
    total_balance = {}
    for balance in balance_list:
      for each_balance in balance.each_balance:
        currency = each_balance.currency
        if currency in total_balance:
          total_balance[currency].total += each_balance.total
          total_balance[currency].available += each_balance.available
          total_balance[currency].hold += each_balance.hold
        else:
          total_balance[currency] = each_balance
    update_usdt.msg = AccountBalance(
        market_type='Futures',
        exchange='Huobi',
        each_balance=[balance for _, balance in total_balance.items()])
    return update_usdt

  def query_account_position_impl(self):
    update_usdt = self.nprivc.query_usdt_position_info()
    update_cross_usdt = self.nprivc.query_cross_usdt_position_info()
    assert (update_usdt.msg['status'] == 'ok' or
            update_usdt.msg.get('err_code') == 1010), update_usdt
    assert (update_cross_usdt.msg['status'] == 'ok' or
            update_cross_usdt.msg.get('err_code') == 1010), update_cross_usdt
    update_usdt.msg['data'] += update_cross_usdt.msg['data']

    account_position = AccountPosition(market_type='Futures', exchange='Huobi')
    linear_msg = self._linear_parser.parse_account_position(json.dumps(update_usdt.msg))
    for msg in (linear_msg,):
      position = AccountPosition()
      json_format.Parse(msg, position)
      assert position.HasField('market_type')
      assert position.HasField('exchange')
      account_position.MergeFrom(position)
    update_usdt.msg = account_position
    return update_usdt

  def query_fills_impl(self, product):
    raise NotImplementedError()

  def query_list_orders_impl(self, product):
    raise NotImplementedError()

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()


class HuobiSwapFuturesPrivateClient(privcb.PrivateClientBase):
  ProductType = HuobiFuturesProduct

  def __init__(self, key_file):
    privcb.PrivateClientBase.__init__(self)
    self.nprivc = nprivc.HuobiSwapNativePrivateClient(key_file=key_file)
    inverse_swap_pi_bundle = huobi_futures_swap_product_info._product_infos
    assert inverse_swap_pi_bundle.HasField('mea')
    self._inverse_parser = PyOrderParser.from_mea_str(inverse_swap_pi_bundle.mea)
    inverse_swap_oss_config = OrderSystemConfig(
      mea=inverse_swap_pi_bundle.mea,
      products=ProductConfig(norms=[pi.symbol for pi in inverse_swap_pi_bundle.product_infos]))
    self._inverse_parser.add_product(inverse_swap_oss_config.SerializeToString())

  def query_account_balance_impl(self):
    update_non_usdt = self.nprivc.query_account_info()
    msg_non_usdt = self._inverse_parser.parse_account_balance(
        json.dumps(update_non_usdt.msg))

    balance_list = []
    for msg in (msg_non_usdt,):
      account_balance = AccountBalance()
      json_format.Parse(msg, account_balance)
      assert account_balance.HasField('market_type')
      assert account_balance.HasField('exchange')
      balance_list.append(account_balance)
    total_balance = {}
    for balance in balance_list:
      for each_balance in balance.each_balance:
        currency = each_balance.currency
        if currency in total_balance:
          total_balance[currency].total += each_balance.total
          total_balance[currency].available += each_balance.available
          total_balance[currency].hold += each_balance.hold
        else:
          total_balance[currency] = each_balance
    update_non_usdt.msg = AccountBalance(
        market_type='Futures',
        exchange='Huobi',
        each_balance=[balance for _, balance in total_balance.items()])
    return update_non_usdt

  def query_account_position_impl(self):
    update = self.nprivc.query_position_info()
    account_position = AccountPosition(market_type='Futures', exchange='Huobi')
    inverse_msg = self._inverse_parser.parse_account_position(json.dumps(update.msg))
    for msg in (inverse_msg,):
      position = AccountPosition()
      json_format.Parse(msg, position)
      assert position.HasField('market_type')
      assert position.HasField('exchange')
      account_position.MergeFrom(position)
    update.msg = account_position
    return update

  def query_fills_impl(self, product):
    raise NotImplementedError()

  def query_list_orders_impl(self, product):
    raise NotImplementedError()

  def query_submit_impl(self, product, side, price, qty):
    raise NotImplementedError()

  def query_cancel_product_impl(self, product):
    raise NotImplementedError()



if __name__ == "__main__":
  client = HuobiSwapFuturesPrivateClient(
      key_file='../../coin_key/view/prestoinvt05_swap/Huobi/view_key.json')
  print(client.query_account_balance().msg)
  print(client.query_account_position().msg)
